#include "common.cuh"
#include "utilities.cuh"
#include "inversion.cuh"
#include "blocklu.cuh"
#include "lutrtri.cuh"
#include "cpu_blas.h"
#include "gpu_blas.cuh"
#define INIT_N 12800
// #define TIME_LOGGER
//#define VERIFY_RESULT
#ifdef _WIN32
#define EXPORT_SYMBOL __declspec(dllexport)
#else
#define EXPORT_SYMBOL 
#endif

// #define VERIFY_RESULT
//#define TIME_LOGGER

#ifdef VERIFY_RESULT
#include <mkl.h>
#endif 

#ifdef TIME_LOGGER
#include <chrono>
using namespace std::chrono;
#endif

COMPUTE_STATE get_next_state(COMPUTE_STATE cur_state, OPERATION op) {
    switch(cur_state) {
        case INITAL: {
            switch (op)
            {
                case GEMV: return INITAL;
                case GETRF: return A33_MUL_X3;
                default: {assert(false); break;}
            }
        }
        case A33_MUL_X3: {
            switch (op)
            {
                case GEMV: return A13_MUL_X3;
                case GETRF: return A33_MUL_X3;
                case GETRS: return INITAL;
                default: {assert(false); break;}
            }
        }
        case A13_MUL_X3: {
            assert(op == GEMV);
            return A23_MUL_X3;
        }
        case A23_MUL_X3: {
            assert(op == GEMV);
            return SOLVE_X1X2;
        }
        case SOLVE_X1X2: {
            assert(op == GETRS);
            return A31_MUL_X1;
        }
        case A31_MUL_X1: {
            assert(op == GEMV);
            return A32_MUL_X2;
        }
        case A32_MUL_X2: {
            assert(op == GEMV);
            return A33_MUL_X3;
        }
    }
    return INITAL;
}


namespace CPU {
Global_Info global_info;

#ifdef VERIFY_RESULT
int verify_initialzed = 0;
double* y_verify;
double* rhs_bak;
double* matrix_bak;
double max_gemv_res;
double max_getrs_res;
int getrf_step;
#endif 

#ifdef TIME_LOGGER
int gemv_counter = 0;
int getrf_counter = 0;
int getrs_counter = 0;
double gemv_sum_tim = 0;
//double getrf_sum_tim = 0;
double getrs_sum_tim = 0;
#endif

void init_global_info() {
    global_info.init = 1;
    global_info.matrix =  (double*) malloc (sizeof(double)*INIT_N*INIT_N);
    global_info.inverse =  (double*) malloc (sizeof(double)*INIT_N*INIT_N);
    global_info.fractor =  global_info.inverse;
    global_info.workplace =  (double*) malloc (sizeof(double)*INIT_N*INIT_N / 16);
    global_info.rhs = (double*) malloc (sizeof(double)*INIT_N);
    global_info.x = (double*) malloc (sizeof(double)*INIT_N);
    global_info.scale = (double*) malloc (sizeof(double)*INIT_N);
    global_info.ipiv = (lapack_int*) malloc(sizeof(lapack_int)*INIT_N);
    global_info.n = 0;
    global_info.last_n = 0;
    global_info.lda = INIT_N;
    global_info.compute_state = INITAL;
    global_info.compute_mode = GPU_ITERATION;
    //initalize ipiv
    initalize_ipiv(global_info.ipiv, global_info.lda);
}
};

namespace DEVICE {
SOLVER_HANDLE solverH;
FACTOR_HANDLE factorH;
MATRIX_SET matrix_set;
VECTOR_SET vector_set;

void init_matrix(MATRIX* mat, size_t max_n) {
    CUDA_RT_CALL(cudaMalloc ((void**)&mat->data, sizeof(double)* max_n * max_n));
    mat->lda = max_n;
    mat->n = 0;
}

void init_solveH() {
    solverH.init = 1;
    CUDA_RT_CALL(cudaMallocHost ((void**)&solverH.pinned_mem, sizeof(double)*INIT_N*INIT_N / 8));
    CUDA_RT_CALL(cusolverDnCreate(&solverH.cudenseH));
    CUDA_RT_CALL(cublasCreate(&solverH.cublasH));
    CUDA_RT_CALL(cudaMalloc ((void**)&solverH.workspace , sizeof(double)*INIT_N*INIT_N) / 8);
    CUDA_RT_CALL(cudaMalloc ((void**)&solverH.devInfo , sizeof(int)));
    for(int i = 0; i < 2; i++) {
        CUDA_RT_CALL(cublasCreate(&solverH.streamH[i]));
        CUDA_RT_CALL(cudaStreamCreate(&solverH.streams[i]));
        CUDA_RT_CALL(cublasSetStream(solverH.streamH[i], solverH.streams[i]));
        CUDA_RT_CALL(cudaEventCreate(&solverH.events[i]));
    }
    CUDA_RT_CALL(cudaDeviceSynchronize());
}

void init_factorH() {
    CUDA_RT_CALL(cudaMalloc ((void**)&factorH.A , sizeof(double)*INIT_N*INIT_N));
    CUDA_RT_CALL(cudaMalloc ((void**)&factorH.inverse , sizeof(double)*INIT_N*INIT_N));
    CUDA_RT_CALL(cudaMalloc ((void**)&factorH.x , sizeof(double)*INIT_N));
    CUDA_RT_CALL(cudaMalloc ((void**)&factorH.rhs , sizeof(double)*INIT_N));
    CUDA_RT_CALL(cudaMalloc ((void**)&factorH.ipiv , sizeof(int)*INIT_N));
    initalize_ipiv(factorH.ipiv, INIT_N);
    CUDA_RT_CALL(cudaDeviceSynchronize());
    factorH.fractor = factorH.inverse;
    factorH.last_n = 0;
};

void init_matrix_set() {
    init_matrix(&matrix_set.A13, INIT_N / 2);
    init_matrix(&matrix_set.A23, INIT_N / 2);
    init_matrix(&matrix_set.A31, INIT_N / 2);
    init_matrix(&matrix_set.A32, INIT_N / 2);
    init_matrix(&matrix_set.A33, INIT_N);
    CUDA_RT_CALL(cudaDeviceSynchronize());
}

void init_vector_set() {
    int length = (INIT_N / 2) * 3;
    CUDA_RT_CALL(cudaMalloc ((void**)&vector_set.x1 , 3 * sizeof(double)*length));
    CUDA_RT_CALL(cudaDeviceSynchronize());
    vector_set.x2 = vector_set.x1 + length;
    vector_set.x3 = vector_set.x1 + 2*length;
}

void update_matrix(SOLVER_HANDLE handle, const double* src, double* dst, int m, int n, int lda, int ldb) {
    double* A12 = handle.pinned_mem;
    double* A2 = handle.pinned_mem + (n - m) * m;
    CPU::memcpy_submatrix(&src[m * lda], A12, m, n - m, lda, m);
    CPU::memcpy_submatrix(&src[m], A2, n - m, n, lda, n - m);

    setMatrix(A12, &dst[m * ldb], m, n - m, m, ldb);
    setMatrix(A2, &dst[m], n - m, n, n - m, ldb);
}

void ddm_gemv(COMPUTE_STATE state, SOLVER_HANDLE handle, double* result, const double* matrix, const double* x, int n, int lda) {
    MATRIX* d_matrix = NULL; double* d_x = handle.workspace; double* d_y = handle.workspace + n; 
    switch(state) {
        case INITAL: { d_matrix = &matrix_set.A33; break; }
        case A33_MUL_X3: { d_matrix = &matrix_set.A33; break; }
        case A13_MUL_X3: { d_matrix = &matrix_set.A13; break; }
        case A23_MUL_X3: { d_matrix = &matrix_set.A23; break; }
        case SOLVE_X1X2: { assert(false); break; }
        case A31_MUL_X1: { d_matrix = &matrix_set.A31; break; }
        case A32_MUL_X2: { d_matrix = &matrix_set.A32; break; }
    }
    switch(state) {
        case INITAL:     {setMatrix(x, d_x, n, 1, n, n); break;}
        case A33_MUL_X3: { d_x = vector_set.x3; setMatrix(x, d_x, n, 1, n, n); break; }
        case A13_MUL_X3: { d_x = vector_set.x3; break; }
        case A23_MUL_X3: { d_x = vector_set.x3; break; }
        case SOLVE_X1X2: { assert(false); break; }
        case A31_MUL_X1: { d_x = vector_set.x1; break; }
        case A32_MUL_X2: { d_x = vector_set.x2; break; }
    }

    
    if(d_matrix->n == 0 || state == INITAL) {
        setMatrix(matrix, d_matrix->data, n, n, lda, d_matrix->lda);
    }
    else if(d_matrix->n < n) {
        update_matrix(solverH, matrix, d_matrix->data, d_matrix->n, n, lda, d_matrix->lda);
    }
    d_matrix->n = n;
    
    DEVICE::gemv(handle.cublasH, d_y, d_matrix->data, d_x, n, n, d_matrix->lda, 1.0, 0.0);
    getMatrix(d_y, result, n, 1, n, n);
}

};

void fractorize(double* matrix, int n) {
    //initalize global info
    if(CPU::global_info.n < n) {
        CPU::global_info.last_n =  CPU::global_info.n;
        CPU::global_info.n = n;
    }
    if(CPU::global_info.last_n == 0) {
        CPU::merge_Amij_to_A(CPU::global_info.matrix, 
                            &matrix[0 + 0 * n], &matrix[0 + (n/2) * n], &matrix[(n/2) + 0*n], &matrix[(n/2) + (n/2)*n],
                            (n/2), (n/2), CPU::global_info.lda, n);
        //extract the diagonal elements
        CPU::extract_scale(CPU::global_info.matrix, CPU::global_info.scale, n, CPU::global_info.lda);
        //scale matrix and rhs
        CPU::scale_matrix(CPU::global_info.matrix, CPU::global_info.scale, n, n, CPU::global_info.lda);
    } else {
        {
            using namespace CPU;
            // write_matrix(matrix, n, n, n, "../debug/Am.bin");
            int offset = global_info.last_n / 2;
            int size = (global_info.n - global_info.last_n) / 2;
            //merge new crack elements action to A
            merge_Amij_to_A(&global_info.matrix[(2*offset) * global_info.lda], 
                            &matrix[offset * n], 
                            &matrix[(2*offset + size) * n],
                            &matrix[offset * n + (offset + size)],
                            &matrix[(2*offset + size) * n + (offset + size)],
                            (size + offset), size, global_info.lda, n);

            //merge existing element action on new ones
            merge_Amij_to_A(&global_info.matrix[(2*offset)], 
                            &matrix[offset], 
                            &matrix[(offset + size) * n + offset],
                            &matrix[(2*offset + size)],
                            &matrix[(offset + size) * n + (2*offset + size)],
                            size, offset, global_info.lda, n);

            //extract new diagonals
            extract_scale(&global_info.matrix[(2*offset) * global_info.lda + (2*offset)], 
                        &global_info.scale[global_info.last_n], 
                        n - global_info.last_n, global_info.lda);
            //scale A12 with old diagonals
            scale_matrix(&global_info.matrix[(2*offset) * global_info.lda],
                        &global_info.scale[0],
                        2*offset, 2*size, global_info.lda);
            //scale A2 with new diagonals
            scale_matrix(&global_info.matrix[2*offset],
                        &global_info.scale[2*offset],
                        2*size, n, global_info.lda);
            // write_matrix(global_info.scale, n, 1, global_info.lda, "../debug/A_diag.bin");
            // write_matrix(global_info.matrix, n, n, global_info.lda, "../debug/A.bin");
        }
    }
    //update inversion state, if global matrix has no inverse, then perform update
    if(CPU::global_info.last_n == 0) {
        // -------------------------------------- INVERSION --------------------------------------
        //copy matrix to inverse
        CPU::memcpy_submatrix(CPU::global_info.matrix, CPU::global_info.inverse, n, n, CPU::global_info.lda, CPU::global_info.lda);
        //perform inversion in-place
        CPU::inplace_inversion(CPU::global_info.inverse, CPU::global_info.ipiv, n, CPU::global_info.lda);
        CPU::global_info.last_n = n;
    } 

    if(CPU::global_info.compute_mode == CPU_ITERATION) {
        if(CPU::global_info.last_n < n) {
            CPU::update_inversion(CPU::global_info.matrix, CPU::global_info.inverse, CPU::global_info.ipiv, CPU::global_info.workplace, 
                            CPU::global_info.last_n, n, CPU::global_info.lda, CPU::global_info.lda);
            CPU::global_info.last_n = n;
        }
    } else if(CPU::global_info.compute_mode == GPU_ITERATION) {
        if(DEVICE::factorH.last_n == 0) {
            DEVICE::setMatrix(CPU::global_info.inverse, DEVICE::factorH.inverse, CPU::global_info.last_n, CPU::global_info.last_n, CPU::global_info.lda, CPU::global_info.lda);
            DEVICE::setMatrix(CPU::global_info.matrix, DEVICE::factorH.A, n, n, CPU::global_info.lda, CPU::global_info.lda);
            DEVICE::factorH.last_n = CPU::global_info.last_n;
        } else {
            DEVICE::update_matrix(DEVICE::solverH, CPU::global_info.matrix, DEVICE::factorH.A, DEVICE::factorH.last_n, n, CPU::global_info.lda, CPU::global_info.lda);
        }
        if(DEVICE::factorH.last_n < n) {
            DEVICE::update_inversion(DEVICE::solverH, DEVICE::factorH.A, DEVICE::factorH.inverse, DEVICE::factorH.ipiv,
                                    DEVICE::factorH.last_n, n, CPU::global_info.lda, CPU::global_info.lda);
            DEVICE::factorH.last_n = n;
        }
    } else {
        assert(false);
    }
}

void solve_rhs(double* matrix, double* rhs, int n) {
    CPU::merge_bi_to_b(CPU::global_info.rhs, &rhs[0], &rhs[(n/2)], (n/2));
    CPU::scale_matrix(CPU::global_info.rhs, CPU::global_info.scale, n, 1, n);
    // -------------------------------------- INVERSION --------------------------------------
    if(CPU::global_info.compute_mode == CPU_ITERATION) {
        CPU::gemv(CPU::global_info.x, CPU::global_info.inverse, CPU::global_info.rhs, n, n, CPU::global_info.lda, 1.0, 0.0);
        CPU::scatter_x_to_xi(&rhs[0], &rhs[(n/2)], CPU::global_info.x, (n/2));
    } else if(CPU::global_info.compute_mode == GPU_ITERATION){
        DEVICE::setMatrix(CPU::global_info.rhs, DEVICE::factorH.rhs, n, 1, n, n);
        DEVICE::gemv(DEVICE::solverH.cublasH, DEVICE::factorH.x, DEVICE::factorH.inverse, DEVICE::factorH.rhs, n, n, CPU::global_info.lda, 1.0, 0.0);
        DEVICE::scatter_x_to_xi(DEVICE::vector_set.x1, DEVICE::vector_set.x2, DEVICE::factorH.x, n / 2);
        DEVICE::getMatrix(DEVICE::vector_set.x1, &rhs[0], n / 2, 1, n / 2, n / 2);
        DEVICE::getMatrix(DEVICE::vector_set.x2, &rhs[(n/2)], n / 2, 1, n / 2, n / 2);
    } else {
        assert(false);
    }
}

#ifdef VERIFY_RESULT
void verifyInit() {
    if(!CPU::verify_initialzed) {
        printf("init for verify result\n");
        CPU::verify_initialzed = 1;
        CPU::y_verify = (double*) malloc(sizeof(double) * INIT_N);
        CPU::rhs_bak = (double*) malloc(sizeof(double) * INIT_N);
        CPU::matrix_bak = (double*) malloc(sizeof(double) * INIT_N * INIT_N);
        CPU::max_gemv_res = 0;
        CPU::max_getrs_res = 0;
        CPU::getrf_step = 0;
    }
}
void updateRes(int type, double res) {
    if(type == 0) { // gemv
        CPU::max_gemv_res = max(CPU::max_gemv_res, res);
    }
    else { // getrs
        CPU::max_getrs_res = max(CPU::max_getrs_res, res);
    }
}
void printRes(int n) {
    CPU::getrf_step++;
    printf("getrf #%d, size : %d, max gemv res : %.e, max getrs res : %.e\n", CPU::getrf_step, n, CPU::max_gemv_res, CPU::max_getrs_res);
}
#endif

#ifdef TIME_LOGGER
void updateTimer(int type, double tim, int n) {
    if(type == 0) { // gemv
        CPU::gemv_counter++;
        CPU::gemv_sum_tim += tim;
    }
    else if(type == 1) { // getrf
        CPU::getrf_counter++;
        printf("getrf #%d, size : %d, getrf time : %lf, gemv count : %d, gemv time : %lf, getrs count : %d, getrs time : %lf\n", 
            CPU::getrf_counter, n, tim, CPU::gemv_counter, CPU::gemv_sum_tim, CPU::getrs_counter, CPU::getrs_sum_tim);
        CPU::gemv_counter = 0;
        CPU::gemv_sum_tim = 0;
        CPU::getrs_counter = 0;
        CPU::getrs_sum_tim = 0;
    }
    else { // getrs
        CPU::getrs_counter++;
        CPU::getrs_sum_tim += tim;
    }
}
#endif

extern "C" EXPORT_SYMBOL void cxxgemv(double* y, double* matrix, double* x_full, int x_start, int n, int lda) {

#ifdef TIME_LOGGER
    auto start = steady_clock::now();
#endif
    
    TIC("cxxgemv");
    double *x = x_full + x_start; 
    if(CPU::global_info.init == 1 && CPU::global_info.compute_mode == GPU_ITERATION) {
        DEVICE::ddm_gemv(CPU::global_info.compute_state, DEVICE::solverH, y, matrix, x, n, lda);
    } else {
        CPU::gemv(y, matrix, x, n, n, lda, 1.0, 0.0);
    }
    
#ifdef VERIFY_RESULT
    verifyInit();
    cblas_dgemv (CblasColMajor, // layout
        CblasNoTrans, // trans
        n, // m
        n, // n
        1.0, // alpha
        matrix, // a
        lda, // lda
        x, // x
        1, // incx
        0.0, // beta
        CPU::y_verify, // y
        1 ); // incy

    double norm_ref = 0.0; 
    double norm_ans = 0.0;
    for(int i = 0; i < n; i++) {
        //CPU::y_verify[i] *= 1.1;
        norm_ref += CPU::y_verify[i] * CPU::y_verify[i];
        norm_ans += (CPU::y_verify[i] - y[i]) * (CPU::y_verify[i] - y[i]);
    }
    //updateRes(0, norm_ans / norm_ref);
    updateRes(0, sqrt(norm_ans) / (sqrt(norm_ref) + 1e-5));
#endif

    //update computation state
    CPU::global_info.compute_state = get_next_state(CPU::global_info.compute_state, GEMV);
        TOC();

#ifdef TIME_LOGGER
    auto end = steady_clock::now();
    duration<double> elapsed_seconds = end - start;
    //printf("gemv took %lf seconds.\n", elapsed_seconds.count());
    updateTimer(0, elapsed_seconds.count(), n);
#endif

}

extern "C" EXPORT_SYMBOL void cxxgetrf(double* matrix, int n) {

#ifdef TIME_LOGGER
    auto start = steady_clock::now();
#endif

#ifdef VERIFY_RESULT
    verifyInit();
    printRes(n);
        memcpy(CPU::matrix_bak, matrix, sizeof(double) * n * n);
#endif

    TIC("cxxgetrf");
    if(CPU::global_info.init == 0) {
        CPU::init_global_info();
        DEVICE::init_solveH();
        DEVICE::init_factorH();
        DEVICE::init_matrix_set();
        DEVICE::init_vector_set();
        printf("initalize complete\n");
    }

    // if(n < 50) {
    //     if(CPU::global_info.compute_mode != DIRECT)
    //         printf("switch to DIRECT, n = %d\n", n);
    //     CPU::global_info.compute_mode = DIRECT;
    // } else if (n < 500) {
    //     if(CPU::global_info.compute_mode != CPU_ITERATION)
    //         printf("switch to CPU_ITERATION, n = %d\n", n);
    //     CPU::global_info.compute_mode = CPU_ITERATION;
    // } else {
    //     if(CPU::global_info.compute_mode != GPU_ITERATION)
    //         printf("switch to GPU_ITERATION, n = %d\n", n);
    //     CPU::global_info.compute_mode = GPU_ITERATION;
    // }

    if(CPU::global_info.compute_mode == DIRECT) {
        CPU::inplace_inversion(matrix, CPU::global_info.ipiv, n, n);
    } else {
        fractorize(matrix, n);
    }

    //update computation state
    CPU::global_info.compute_state = get_next_state(CPU::global_info.compute_state, GETRF);
        TOC();

#ifdef TIME_LOGGER
    auto end = steady_clock::now();
    duration<double> elapsed_seconds = end - start;
    updateTimer(1, elapsed_seconds.count(), n);
#endif

}

extern "C" EXPORT_SYMBOL void cxxgetrs(double* matrix, double* rhs, int n) {

#ifdef TIME_LOGGER
    auto start = steady_clock::now();
#endif

#ifdef VERIFY_RESULT
    verifyInit();
    memcpy(CPU::rhs_bak, rhs, sizeof(double) * n);
#endif
    
    TIC("cxxgetrs");
    if(CPU::global_info.compute_mode == DIRECT) {
        double* y = CPU::global_info.workplace;
        CPU::gemv(y, matrix, rhs, n, n, n, 1.0, 0.0);
        CPU::memcpy_submatrix(y, rhs, n, 1, n, n);
    } else {
        solve_rhs(matrix, rhs, n);
    }

#ifdef VERIFY_RESULT
    cblas_dgemv (CblasColMajor, // layout
        CblasNoTrans, // trans
        n, // m
        n, // n
        1.0, // alpha
        CPU::matrix_bak, // a
        n, // lda
        rhs, // x
        1, // incx
        0.0, // beta
        CPU::y_verify, // y
        1 ); // incy

    double norm_ref = 0.0; 
    double norm_ans = 0.0;
    for(int i = 0; i < n; i++) {
        //CPU::y_verify[i]++;
        norm_ref += CPU::y_verify[i] * CPU::y_verify[i];
        norm_ans += (CPU::y_verify[i] - CPU::rhs_bak[i]) * (CPU::y_verify[i] - CPU::rhs_bak[i]);
    }
    updateRes(2, sqrt(norm_ans) / (sqrt(norm_ref) + 1e-5));
    #endif

    //update computation state
    CPU::global_info.compute_state = get_next_state(CPU::global_info.compute_state, GETRS);
    
    TOC();

#ifdef TIME_LOGGER
    auto end = steady_clock::now();
    duration<double> elapsed_seconds = end - start;
    //printf("gemv took %lf seconds.\n", elapsed_seconds.count());
    updateTimer(2, elapsed_seconds.count(), n);
#endif

}


extern "C" void ddm_solver(double* matrix, double* rhs, int n) {
    TIC("ddm_solver");
    cxxgetrf(matrix, n);
    cxxgetrs(matrix, rhs, n);
    // fractorize(matrix, n);
    // solve_rhs(matrix, rhs, n);
    TOC();
}
